package cn.websocket.filter;

import cn.websocket.pojo.config.RateLimiterProperties;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;


/**
 * <h3>WebSocket</h3>
 * <p>
 *     ip限流防止脚本，频繁刷新验证码接口
 * </p>
 *
 * @author gdw
 * @since 2025-08-20 11:06:23
 */
@Log4j2
@Component
public class RateLimiterFilter implements Filter {
    @Autowired
    private  final RateLimiterProperties properties;

    private final Cache<String, Bucket> bucketCache = CacheBuilder.newBuilder()
            .expireAfterAccess(10, TimeUnit.MINUTES)
            .maximumSize(10000)
            .build();

    @Autowired
    public RateLimiterFilter(RateLimiterProperties properties) {
        this.properties = properties;
    }

    private Bucket createNewBucket(RateLimiterProperties.Rule rule) {
        Duration duration = parseDuration(rule.getDuration());
        Refill refill = Refill.intervally(rule.getRefill(), duration);
        Bandwidth limit = Bandwidth.classic(rule.getCapacity(), refill);
        return Bucket.builder().addLimit(limit).build();
    }

    private Duration parseDuration(String str) {
        if (str.endsWith("s")) {
            return Duration.ofSeconds(Long.parseLong(str.replace("s", "")));
        } else if (str.endsWith("m")) {
            return Duration.ofMinutes(Long.parseLong(str.replace("m", "")));
        } else if (str.endsWith("h")) {
            return Duration.ofHours(Long.parseLong(str.replace("h", "")));
        } else {
            throw new IllegalArgumentException("Unsupported duration format: " + str);
        }
    }

    private String getClientIP(HttpServletRequest request) {
        String xfHeader = request.getHeader("X-Forwarded-For");
        if (xfHeader != null && !xfHeader.isEmpty() && !"unknown".equalsIgnoreCase(xfHeader)) {
            return xfHeader.split(",")[0].trim();
        }
        String realIp = request.getHeader("X-Real-IP");
        if (realIp != null && !realIp.isEmpty() && !"unknown".equalsIgnoreCase(realIp)) {
            return realIp;
        }
        return request.getRemoteAddr();
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {

    }

    @Override
    public void doFilter(ServletRequest req, ServletResponse res, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest request = (HttpServletRequest) req;
        HttpServletResponse response = (HttpServletResponse) res;

        String path = request.getRequestURI();

        // 找到匹配的限流规则
        Optional<RateLimiterProperties.Rule> matchedRule =
                properties.getRules().stream().filter(rule -> path.startsWith(rule.getPath())).findFirst();

        if (matchedRule.isPresent()) {
            String ip = getClientIP(request);
            String cacheKey = ip + ":" + matchedRule.get().getPath(); // IP+路径 做 key

            Bucket bucket;
            try {
                bucket = bucketCache.get(cacheKey, () -> createNewBucket(matchedRule.get()));
            } catch (ExecutionException e) {
                throw new RuntimeException(e);
            }

            if (!bucket.tryConsume(1)) {
                response.setStatus(429);
                response.setContentType("application/json;charset=UTF-8");
                response.getWriter().write("{\"code\":429, \"msg\":\"请求过于频繁，请稍后再试。\"}");
                return;
            }
            log.info("Path={}, IP={}, tryConsume={}", path, ip, bucket.tryConsume(1));
        }

        chain.doFilter(req, res);
    }

    @Override
    public void destroy() {

    }
}
